#include "StdAfx.h"
#include "InstallationHelper.h"
#include "MultiLinePropertyHandler.h"

CInstallationHelper::CInstallationHelper(VsVersion eVsVersion)
{
	LPCTSTR lpDteClsid = NULL;
	switch (eVsVersion)
	{
	case vsV8:
		m_strLibFileName = _T("cputcrt8.lib");
		lpDteClsid = _T("CLSID\\{BA018599-1DB3-44f9-83B4-461454C84BF8}\\LocalServer32");
		break;
	case vsV9:
		m_strLibFileName = _T("cputcrt9.lib");
		lpDteClsid = _T("CLSID\\{1A5AC6AE-7B95-478C-B422-0E994FD727D6}\\LocalServer32");
		break;
	};

	CPath csDevEnvPath;
	CRegKey csDteClsid;
	ULONG ul = MAX_PATH;
	LPTSTR lp = csDevEnvPath.m_strPath.GetBuffer(MAX_PATH);
	if (ERROR_SUCCESS == csDteClsid.Open(HKEY_CLASSES_ROOT, lpDteClsid, KEY_READ))
	{
		csDteClsid.QueryStringValue(NULL, lp, &ul);
	}
	csDevEnvPath.m_strPath.ReleaseBuffer();
	csDevEnvPath.UnquoteSpaces();
	if (csDevEnvPath.FileExists())
	{
		csDevEnvPath.RemoveFileSpec();
		m_strUserTypeFilePath = csDevEnvPath.m_strPath;
		m_strUserTypeFilePath.Append(_T("\\UserType.dat"));
		csDevEnvPath.RemoveFileSpec();
		csDevEnvPath.RemoveFileSpec();
		m_strStudioRootFolder = csDevEnvPath.m_strPath;
	}
}

CInstallationHelper::~CInstallationHelper(void)
{
}

void CInstallationHelper::RegisterUtUserType()
{
	CAtlStringW wstr;
	BOOL bIsUnicode = FALSE;
	if (ReadUserTypeFile(wstr, bIsUnicode))
	{
		CMultiLinePropertyHandler csParser;
		csParser.Parse(wstr, L"\r\n");
		csParser.Add(L"ut_private");
		csParser.Add(L"ut_protected");
		if (csParser.IsModified())
		{
			csParser.GetPropertyLine(wstr, L"\r\n");
			WriteUserTypeFile(wstr, bIsUnicode);
		} 

	} else
	{
		WriteUserTypeFile(L"ut_private\r\nut_protected\r\n", FALSE);
	}
}

void CInstallationHelper::UnregisterUtUserType()
{
	CAtlStringW wstr;
	BOOL bIsUnicode = FALSE;
	if (ReadUserTypeFile(wstr, bIsUnicode))
	{
		CMultiLinePropertyHandler csParser;
		csParser.Parse(wstr, L"\r\n");
		csParser.RemoveIfExist(L"ut_private");
		csParser.RemoveIfExist(L"ut_protected");
		if (csParser.IsModified())
		{
			csParser.GetPropertyLine(wstr, L"\r\n");
			if (wstr.GetLength() < 2)
			{
				DeleteFile(m_strUserTypeFilePath);
			} else
			{
				WriteUserTypeFile(wstr, bIsUnicode);
			}
		}
	}
}

BOOL CInstallationHelper::ReadUserTypeFile(OUT CAtlStringW & wstr, OUT BOOL & bIsUnicode)
{
	BOOL bRet = FALSE;
	WORD w = 0;
	CAtlFile csUserTypeDat;
	ULARGE_INTEGER liSize;
	if (SUCCEEDED(csUserTypeDat.Create(m_strUserTypeFilePath, GENERIC_READ, 0, OPEN_EXISTING)) &&
		SUCCEEDED(csUserTypeDat.GetSize(liSize.QuadPart)) &&
		(liSize.HighPart == 0) &&
		(liSize.QuadPart > 2) &&
		SUCCEEDED(csUserTypeDat.Read(&w, sizeof(w))))
	{
		CAtlStringW wstrBuf;
		if (w == 0xFEFF)
		{
			int nLen = (liSize.LowPart - sizeof(w))/sizeof(wchar_t);
			bRet = SUCCEEDED(csUserTypeDat.Read(wstr.GetBuffer(nLen), liSize.LowPart - sizeof(w)));
			wstr.ReleaseBuffer(nLen);
			bIsUnicode = TRUE;
		} else
		{
			CAtlStringA str;
			csUserTypeDat.Seek(0, FILE_BEGIN);
			bRet = SUCCEEDED(csUserTypeDat.Read(str.GetBuffer(liSize.LowPart), liSize.LowPart));
			str.ReleaseBuffer(liSize.LowPart);
			wstr = str;
			bIsUnicode = FALSE;
		}
	}
	return bRet;
}

BOOL CInstallationHelper::WriteUserTypeFile(IN const CAtlStringW & wstr, IN BOOL bIsUnicode)
{
	BOOL bRet = FALSE;
	CAtlFile csUserTypeDat;
	if (SUCCEEDED(csUserTypeDat.Create(m_strUserTypeFilePath, GENERIC_WRITE, 0, CREATE_NEW)) ||
		SUCCEEDED(csUserTypeDat.Create(m_strUserTypeFilePath, GENERIC_WRITE, 0, TRUNCATE_EXISTING)))
	{
		if (bIsUnicode)
		{
			WORD w = 0xFEFF;
			if (SUCCEEDED(csUserTypeDat.Write(&w, sizeof(w))) &&
				SUCCEEDED(csUserTypeDat.Write(wstr.GetString(), wstr.GetLength() * sizeof(wchar_t))))
				bRet = TRUE;
		} else
		{
			CAtlStringA str(wstr.GetString());
			bRet = SUCCEEDED(csUserTypeDat.Write(str.GetString(), str.GetLength()));
		}
	}
	return bRet;
}

void CInstallationHelper::MoveLibFiles()
{
	CPath path;
	GetModuleFileName( _AtlModule.GetResourceInstance(), path.m_strPath.GetBuffer(MAX_PATH), MAX_PATH);
	path.m_strPath.ReleaseBuffer();
	path.RemoveFileSpec();

	CPath pathUtAssert(path);
	pathUtAssert.Append(_T("utassert.h"));
	CPath pathUtTestBase(path);
	pathUtTestBase.Append(_T("uttestbase.h"));
	CPath pathCputLib(path);
	pathCputLib.Append(m_strLibFileName);

	if (pathUtAssert.FileExists() &&
		pathUtTestBase.FileExists() &&
		pathCputLib.FileExists())
	{
		CopyFile(pathUtAssert, m_strStudioRootFolder + _T("\\VC\\include\\utassert.h"), FALSE);
		CopyFile(pathUtTestBase, m_strStudioRootFolder + _T("\\VC\\include\\uttestbase.h"), FALSE);
		CopyFile(pathCputLib, m_strStudioRootFolder + _T("\\VC\\lib\\") + m_strLibFileName, FALSE);
	}
}

void CInstallationHelper::RemoveLibFiles()
{
	DeleteFile(m_strStudioRootFolder + _T("\\VC\\include\\utassert.h"));
	DeleteFile(m_strStudioRootFolder + _T("\\VC\\include\\uttestbase.h"));
	DeleteFile(m_strStudioRootFolder + _T("\\VC\\lib\\") + m_strLibFileName);
}

void CInstallationHelper::Install()
{
#ifdef _VS2008
	CInstallationHelper csDteHelper(vsV9);
#else
	CInstallationHelper csDteHelper(vsV8);
#endif

	csDteHelper.RegisterUtUserType();
	csDteHelper.MoveLibFiles();
}

void CInstallationHelper::Uninstall()
{
#ifdef _VS2008
	CInstallationHelper csDteHelper(vsV9);
#else
	CInstallationHelper csDteHelper(vsV8);
#endif

	csDteHelper.UnregisterUtUserType();
	csDteHelper.RemoveLibFiles();
}
